from typing import List, Union

from bson.objectid import ObjectId

from src.basedata.infrastructure.client.redis.repository import RedisRepository
from src.seedwork.domain.entities import Entity
from src.seedwork.application.exceptions import EntityNotFoundException


class Repository:
    pass


class InMemoryRepository(Repository):
    def __init__(self) -> None:
        self.objects = {}

    def get_by_id(self, id) -> Entity:
        try:
            return self.objects[id]
        except KeyError:
            raise EntityNotFoundException

    def insert(self, entity: Entity):
        assert issubclass(entity.__class__, Entity)
        self.objects[entity.id] = entity

    def update(self, entity: Entity):
        assert issubclass(entity.__class__, Entity)
        self.objects[entity.id] = entity

    def delete(self, entity_id):
        del self.objects[entity_id]


class MongoDBRepository:

    def __init__(self, session, table_name):
        self.session = session
        self.table_name = table_name

    async def aggregate(self, pipelines: list):
        table = self.session.get_collection(self.table_name)
        cursor = table.aggregate(pipelines)
        records = []
        # 结果直接转成字典
        async for item in cursor:
            records.append(item)
        return records

    async def query_record(self, query: dict):
        table = self.session.get_collection(self.table_name)
        record = await table.find_one(query)
        return record

    async def query_records(self, query: dict):
        table = self.session.get_collection(self.table_name)
        cursor = table.find(query)
        records = []
        async for record in cursor:
            records.append(record)
        return records

    async def query_records_limit(self, query: dict, page: int, num: int, sort_param: str, reverse=False):
        page = int(page)
        num = int(num)
        skip_num = num * (page - 1)
        table = self.session.get_collection(self.table_name)
        cursor = table.find(query).sort(sort_param, -1 if reverse else 1).skip(skip_num).limit(num)
        records = []
        async for record in cursor:
            records.append(record)
        return records

    async def query_records_limit_offset(self, query: dict, limit: int, offset: int, sort_param: str, reverse=False):
        table = self.session.get_collection(self.table_name)
        cursor = table.find(query).sort(sort_param, -1 if reverse else 1).skip(offset).limit(limit)
        records = []
        async for record in cursor:
            records.append(record)
        return records

    async def query_count(self, query: dict):
        table = self.session.get_collection(self.table_name)
        count = await table.count_documents(query)
        return count

    async def insert(self, record: dict):
        table = self.session.get_collection(self.table_name)
        result = await table.insert_one(record)
        return str(result.inserted_id)

    async def insert_many(self, records: List[dict]):
        table = self.session.get_collection(self.table_name)
        result = await table.insert_many(records)
        return result

    async def update_one(self, query: dict, value: dict):
        table = self.session.get_collection(self.table_name)
        await table.update_one(query, {"$set": value})

    async def update_many(self, query: dict, value: dict):
        table = self.session.get_collection(self.table_name)
        return await table.update_many(query, {"$set": value})

    async def delete(self, where: dict):
        value = {
            "obsoleted": True,
        }
        await self.update_many(where, value)

    def _convert_to_str(self, m):
        if m:
            for k in m:
                if isinstance(m[k], ObjectId):
                    m[k] = str(m[k])
        return m


class MongoDBMixedRedisRepository(MongoDBRepository, RedisRepository):
    def __init__(self, session, table_name, *args):
        super(MongoDBMixedRedisRepository, self).__init__(session, table_name)
        RedisRepository.__init__(self, *args)
